{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "grBmytrShbUE"
      },
      "source": [
        "# High-performance simulations with TFF\n",
        "\n",
        "This tutorial will describe how to setup high-performance simulations with TFF\n",
        "in a variety of common scenarios.\n",
        "\n",
        "TODO(b/134543154): Populate the content, some of the things to cover here:\n",
        "- using GPUs in a single-machine setup,\n",
        "- multi-machine setup on GCP/GKE, with and without TPUs,\n",
        "- interfacing MapReduce-like backends,\n",
        "- current limitations and when/how they will be relaxed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iPFgLeZIsZ3Q"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/federated/tutorials/simulations\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/federated/blob/master/docs/tutorials/simulations.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/federated/blob/master/docs/tutorials/simulations.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/federated/docs/tutorials/simulations.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yiq_MY4LopET"
      },
      "source": [
        "## Before we begin\n",
        "\n",
        "First, make sure your notebook is connected to a backend that has the relevant\n",
        "components (including gRPC dependencies for multi-machine scenarios) compiled."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_zFenI3IPpgI"
      },
      "source": [
        "Now, let's start by loading the MNIST example from the TFF website, and\n",
        "declaring the Python function that will run a small experiment loop over\n",
        "a group of 10 clients."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "height": 35
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 431,
          "status": "ok",
          "timestamp": 1592938309587,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "ke7EyuvG0Zyn",
        "outputId": "5e577237-a756-40d4-ff81-a2adf85452f2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/bin/sh: pip: command not found\n"
          ]
        }
      ],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install --quiet --upgrade tensorflow-federated-nightly\n",
        "!pip install --quiet --upgrade nest-asyncio\n",
        "\n",
        "import nest_asyncio\n",
        "nest_asyncio.apply()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 23138,
          "status": "ok",
          "timestamp": 1592938332731,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "2dVPgxN0MdG2"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "import time\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "import tensorflow_federated as tff\n",
        "\n",
        "source, _ = tff.simulation.datasets.emnist.load_data()\n",
        "\n",
        "\n",
        "def map_fn(example):\n",
        "  return collections.OrderedDict(\n",
        "      x=tf.reshape(example['pixels'], [-1, 784]), y=example['label'])\n",
        "\n",
        "\n",
        "def client_data(n):\n",
        "  ds = source.create_tf_dataset_for_client(source.client_ids[n])\n",
        "  return ds.repeat(10).shuffle(500).batch(20).map(map_fn)\n",
        "\n",
        "\n",
        "train_data = [client_data(n) for n in range(10)]\n",
        "element_spec = train_data[0].element_spec\n",
        "\n",
        "\n",
        "def model_fn():\n",
        "  model = tf.keras.models.Sequential([\n",
        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
        "      tf.keras.layers.Dense(units=10, kernel_initializer='zeros'),\n",
        "      tf.keras.layers.Softmax(),\n",
        "  ])\n",
        "  return tff.learning.from_keras_model(\n",
        "      model,\n",
        "      input_spec=element_spec,\n",
        "      loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])\n",
        "\n",
        "\n",
        "trainer = tff.learning.build_federated_averaging_process(\n",
        "    model_fn, client_optimizer_fn=lambda: tf.keras.optimizers.SGD(0.02))\n",
        "\n",
        "\n",
        "def evaluate(num_rounds=10):\n",
        "  state = trainer.initialize()\n",
        "  for _ in range(num_rounds):\n",
        "    t1 = time.time()\n",
        "    state, metrics = trainer.next(state, train_data)\n",
        "    t2 = time.time()\n",
        "    print('metrics {m}, round time {t:.2f} seconds'.format(\n",
        "        m=metrics, t=t2 - t1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CDHJF7EIiEy-"
      },
      "source": [
        "## Single-machine simulations\n",
        "\n",
        "Now on by default."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "height": 190
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 24224,
          "status": "ok",
          "timestamp": 1592938356963,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "-V6uCS_BMoR9",
        "outputId": "968c4957-a493-48fc-ea9d-c4ce7fdba468"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "metrics \u003csparse_categorical_accuracy=0.13858024775981903,loss=3.0073554515838623\u003e, round time 3.59 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.1796296238899231,loss=2.749046802520752\u003e, round time 2.29 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.21656379103660583,loss=2.514779567718506\u003e, round time 2.33 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.2637860178947449,loss=2.312587261199951\u003e, round time 2.06 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.3334362208843231,loss=2.068122386932373\u003e, round time 2.00 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.3737654387950897,loss=1.9268712997436523\u003e, round time 2.42 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.4296296238899231,loss=1.7216310501098633\u003e, round time 2.20 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.4655349850654602,loss=1.6489890813827515\u003e, round time 2.18 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.5048353672027588,loss=1.5485210418701172\u003e, round time 2.16 seconds\n",
            "metrics \u003csparse_categorical_accuracy=0.5564814805984497,loss=1.4140453338623047\u003e, round time 2.41 seconds\n"
          ]
        }
      ],
      "source": [
        "evaluate()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bZ171NhcNa3M"
      },
      "source": [
        "## Multi-machine simulations on GCP/GKE, GPUs, TPUs, and beyond...\n",
        "\n",
        "Coming very soon."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "High-performance simulations with TFF"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
